// 这个文件用于测试不同实现的GEMM算法性能以及正确性（允许误差范围1e-5）
// 根据测试需要，可以在main函数中修改矩阵大小、GEMM算法等代码

#include "common.h"
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

// 生成随机数填充矩阵的函数
static void generate_random_matrix(float *matrix, int rows, int cols) {
  srand((unsigned)time(NULL));
  for (int i = 0; i < rows * cols; ++i) {
    matrix[i] = (float)(rand() % 100) / 100.0f; // 生成0到1之间的随机浮点数
  }
}

static float **generate_matrix(int M, int N, int K) {
  float **matrix = (float **)malloc(4 * sizeof(float *));

  matrix[0] = (float *)aligned_alloc(64, M * K * sizeof(float));
  matrix[1] = (float *)aligned_alloc(64, K * N * sizeof(float));
  matrix[2] = (float *)aligned_alloc(64, M * N * sizeof(float));
  matrix[3] = (float *)aligned_alloc(64, M * N * sizeof(float));

  generate_random_matrix(matrix[0], M, K);
  generate_random_matrix(matrix[1], K, N);
  generate_random_matrix(matrix[2], M, N);
  memcpy(matrix[3], matrix[2], M * N * sizeof(float));

  return matrix;
}

int main() {
  srand(time(NULL)); // 初始化随机数种子
  // 测试不同尺寸的矩阵
  for (int size = 200; size <= 2000; size += 77) {
    int M = size, N = size, K = size;
    float ALPHA = (float)(rand() % 100) / 100.0f;

    float **matrixs = generate_matrix(M, N, K);
    float *A = matrixs[0];
    float *B = matrixs[1];
    float *C = matrixs[2];
    float *D = matrixs[3];

    struct timespec start, end;
    long seconds, ns;

    clock_gettime(CLOCK_REALTIME, &start);
    gemm_nn_rvv_pack(M, N, K, ALPHA, A, B, C);
    clock_gettime(CLOCK_REALTIME, &end);
    seconds = end.tv_sec - start.tv_sec;
    ns = end.tv_nsec - start.tv_nsec;
    if (start.tv_nsec > end.tv_nsec) {
      --seconds;
      ns += 1000000000;
    }
    printf("[gemm_nn_rvv_pack] Matrix size %dx%dx%d: %ld.%09ld seconds\n", M, N,
           K, seconds, ns);

    clock_gettime(CLOCK_REALTIME, &start);
    // asm("addi x0,x0,1"); // rvtrace开始记录指令的HINT
    gemm_nn_rvv_MNK(M, N, K, ALPHA, A, B, D);
    // asm("addi x0,x0,2"); // rvtrace记录指令结束的HINT
    clock_gettime(CLOCK_REALTIME, &end);
    seconds = end.tv_sec - start.tv_sec;
    ns = end.tv_nsec - start.tv_nsec;
    if (start.tv_nsec > end.tv_nsec) {
      --seconds;
      ns += 1000000000;
    }
    printf("[gemm_nn_rvv_MNK] Matrix size %dx%dx%d: %ld.%09ld seconds\n", M,
           N, K, seconds, ns);

    for (int i = 0; i < M * N; ++i) {
      if (fabs(D[i] - C[i]) > 1e-5) {
        printf("Error: C[%d] = %f, D[%d] = %f\n", i, C[i], i, D[i]);
        return 1;
      }
    }

    free(A);
    free(B);
    free(C);
    free(D);
    free(matrixs);
  }

  return 0;
}
